import pygame
from src.scenes.start_screen import StartScreen
from src.scenes.game_screen import GameScreen
from src.scenes.end_screen import EndScreen
from src.scenes.success_screen import SuccessScreen  # 导入新的成功界面类

pygame.init()
screen = pygame.display.set_mode((800, 600))
pygame.display.set_caption("坦克大战")


def main():
    clock = pygame.time.Clock()
    state = "start"
    game_screen = GameScreen(screen)

    while True:
        if state == "start":
            state = StartScreen(screen).run()
        elif state == "play":
            result = game_screen.run()
            if result is False:
                state = "end"
            elif result == "success":
                state = "success"
            else:
                state = "play"
        elif state == "end":
            state = EndScreen(screen).run()
        elif state == "success":
            state = SuccessScreen(screen).run()
        clock.tick(60)


if __name__ == "__main__":
    main()